{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-08-07 21:51:10] Try to use the default NATS-Bench (size) path from fast_mode=True and path=None.\n"
     ]
    }
   ],
   "source": [
    "from nats_bench import create\n",
    "from nats_bench.api_utils import time_string\n",
    "import numpy as np\n",
    "\n",
    "# Create the API for size search space\n",
    "api = create(None, 'sss', fast_mode=True, verbose=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2021-08-07 21:51:10] There are 32768 architectures on the size search space\n",
      "{'name': 'infer.shape.tiny', 'channels': '8:8:8:16:40', 'genotype': '|nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_3x3~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|', 'num_classes': 10}\n"
     ]
    }
   ],
   "source": [
    "print('{:} There are {:} architectures on the size search space'.format(time_string(), len(api)))\n",
    "\n",
    "# Obtain the 12-th candidate's configureation on CIFAR-10\n",
    "config = api.get_net_config(12, 'cifar10')\n",
    "print(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xautodl  # import this lib -- \"https://github.com/D-X-Y/AutoDL-Projects\", you can use pip install xautodl\n",
    "from xautodl.models import get_cell_based_tiny_net\n",
    "# create the network, which is the sub-class of torch.nn.Module\n",
    "network = get_cell_based_tiny_net(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DynamicShapeTinyNet(\n",
      "  DynamicShapeTinyNet(C=(8, 8, 8, 16, 40), N=1, L=5)\n",
      "  (stem): Sequential(\n",
      "    (0): Conv2d(3, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "    (1): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (cells): ModuleList(\n",
      "    (0): InferCell(\n",
      "      info :: nodes=4, inC=8, outC=8, [1<-(I0-L0) | 2<-(I0-L1,I1-L2) | 3<-(I0-L3,I1-L4,I2-L5)], |nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_3x3~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|\n",
      "      (layers): ModuleList(\n",
      "        (0): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (1): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (2): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (3): Identity()\n",
      "        (4): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (5): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (1): ResNetBasicblock(\n",
      "      ResNetBasicblock(inC=8, outC=8, stride=2)\n",
      "      (conv_a): ReLUConvBN(\n",
      "        (op): Sequential(\n",
      "          (0): ReLU()\n",
      "          (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "          (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (conv_b): ReLUConvBN(\n",
      "        (op): Sequential(\n",
      "          (0): ReLU()\n",
      "          (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (downsample): Sequential(\n",
      "        (0): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
      "        (1): Conv2d(8, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      )\n",
      "    )\n",
      "    (2): InferCell(\n",
      "      info :: nodes=4, inC=8, outC=8, [1<-(I0-L0) | 2<-(I0-L1,I1-L2) | 3<-(I0-L3,I1-L4,I2-L5)], |nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_3x3~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|\n",
      "      (layers): ModuleList(\n",
      "        (0): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (1): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (2): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (3): Identity()\n",
      "        (4): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (5): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (3): ResNetBasicblock(\n",
      "      ResNetBasicblock(inC=8, outC=16, stride=2)\n",
      "      (conv_a): ReLUConvBN(\n",
      "        (op): Sequential(\n",
      "          (0): ReLU()\n",
      "          (1): Conv2d(8, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "          (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (conv_b): ReLUConvBN(\n",
      "        (op): Sequential(\n",
      "          (0): ReLU()\n",
      "          (1): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "          (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "      )\n",
      "      (downsample): Sequential(\n",
      "        (0): AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
      "        (1): Conv2d(8, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "      )\n",
      "    )\n",
      "    (4): InferCell(\n",
      "      info :: nodes=4, inC=16, outC=40, [1<-(I0-L0) | 2<-(I0-L1,I1-L2) | 3<-(I0-L3,I1-L4,I2-L5)], |nor_conv_3x3~0|+|nor_conv_3x3~0|nor_conv_3x3~1|+|skip_connect~0|nor_conv_3x3~1|nor_conv_3x3~2|\n",
      "      (layers): ModuleList(\n",
      "        (0): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(16, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (1): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(16, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (2): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (3): FactorizedReduce(\n",
      "          C_in=16, C_out=40, stride=1\n",
      "          (relu): ReLU()\n",
      "          (conv): Conv2d(16, 40, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
      "          (bn): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "        )\n",
      "        (4): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "        (5): ReLUConvBN(\n",
      "          (op): Sequential(\n",
      "            (0): ReLU()\n",
      "            (1): Conv2d(40, 40, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "            (2): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "          )\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (lastact): Sequential(\n",
      "    (0): BatchNorm2d(40, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    (1): ReLU(inplace=True)\n",
      "  )\n",
      "  (global_pooling): AdaptiveAvgPool2d(output_size=1)\n",
      "  (classifier): Linear(in_features=40, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(network)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model parameters are 0.067378 MB\n"
     ]
    }
   ],
   "source": [
    "from xautodl.utils import count_parameters_in_MB\n",
    "print('The model parameters are {:} MB'.format(count_parameters_in_MB(network)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
